#!/usr/bin/env python3
"""
ComfyUI Dify Integration - 工作流调试脚本
帮助诊断和测试ComfyUI与Dify集成的问题
"""

import json
import sys
import os
from typing import Dict, Any, Optional

def load_workflow(workflow_path: str) -> Optional[Dict[str, Any]]:
    """加载工作流文件"""
    try:
        with open(workflow_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except FileNotFoundError:
        print(f"❌ 工作流文件未找到: {workflow_path}")
        return None
    except json.JSONDecodeError as e:
        print(f"❌ 工作流JSON格式错误: {e}")
        return None

def analyze_workflow(workflow: Dict[str, Any]) -> Dict[str, Any]:
    """分析工作流结构"""
    analysis = {
        "nodes_count": len(workflow),
        "dify_output_nodes": [],
        "text_input_nodes": [],
        "connections": [],
        "potential_issues": []
    }
    
    # 分析每个节点
    for node_id, node_data in workflow.items():
        class_type = node_data.get("class_type", "Unknown")
        inputs = node_data.get("inputs", {})
        
        # 查找Dify输出节点
        if class_type == "DifyTextOutputNode":
            analysis["dify_output_nodes"].append({
                "id": node_id,
                "inputs": inputs,
                "title": node_data.get("_meta", {}).get("title", "")
            })
            
            # 检查text输入
            text_input = inputs.get("text")
            if isinstance(text_input, list) and len(text_input) == 2:
                # 这是一个连接
                source_node = text_input[0]
                output_index = text_input[1]
                analysis["connections"].append({
                    "from": source_node,
                    "to": node_id,
                    "output_index": output_index,
                    "type": "text"
                })
                
                # 检查源节点是否存在
                if source_node not in workflow:
                    analysis["potential_issues"].append(
                        f"节点 {node_id} 的文本输入连接到不存在的节点 {source_node}"
                    )
        
        # 查找文本输入节点
        if "text" in inputs and isinstance(inputs["text"], str):
            analysis["text_input_nodes"].append({
                "id": node_id,
                "class_type": class_type,
                "text_content": inputs["text"][:100] + "..." if len(inputs["text"]) > 100 else inputs["text"]
            })
    
    return analysis

def simulate_dify_output(workflow: Dict[str, Any]) -> Dict[str, Any]:
    """模拟Dify输出节点的执行"""
    results = {}
    
    for node_id, node_data in workflow.items():
        if node_data.get("class_type") == "DifyTextOutputNode":
            inputs = node_data.get("inputs", {})
            text_input = inputs.get("text", "")
            
            # 如果是连接输入，尝试从源节点获取文本
            if isinstance(text_input, list) and len(text_input) == 2:
                source_node_id = text_input[0]
                source_node = workflow.get(source_node_id)
                
                if source_node:
                    source_inputs = source_node.get("inputs", {})
                    # 尝试获取源节点的文本
                    actual_text = source_inputs.get("text", "")
                else:
                    actual_text = ""
            else:
                actual_text = text_input
            
            # 模拟DifyTextOutputNode的输出
            try:
                # 创建符合Dify格式的输出
                dify_output = {
                    "text": actual_text,
                    "files": [],
                    "json": [{"data": []}]
                }
                
                # 如果包含元数据
                if inputs.get("include_metadata", False):
                    dify_output["metadata"] = {
                        "node_id": node_id,
                        "timestamp": "2024-01-01T00:00:00Z",
                        "text_length": len(actual_text)
                    }
                
                results[node_id] = {
                    "status": "success",
                    "output": dify_output,
                    "formatted_output": json.dumps(dify_output, ensure_ascii=False, indent=2)
                }
                
            except Exception as e:
                results[node_id] = {
                    "status": "error",
                    "error": str(e)
                }
    
    return results

def print_analysis(analysis: Dict[str, Any]) -> None:
    """打印分析结果"""
    print("="*60)
    print("🔍 工作流结构分析")
    print("="*60)
    
    print(f"\n📊 基础信息:")
    print(f"  - 节点总数: {analysis['nodes_count']}")
    print(f"  - Dify输出节点数: {len(analysis['dify_output_nodes'])}")
    print(f"  - 文本输入节点数: {len(analysis['text_input_nodes'])}")
    print(f"  - 连接数: {len(analysis['connections'])}")
    
    if analysis['dify_output_nodes']:
        print(f"\n🎯 Dify输出节点:")
        for i, node in enumerate(analysis['dify_output_nodes'], 1):
            print(f"  {i}. 节点 {node['id']} ({node['title']})")
            print(f"     - 格式化为Dify: {node['inputs'].get('format_for_dify', 'N/A')}")
            print(f"     - 包含元数据: {node['inputs'].get('include_metadata', 'N/A')}")
    
    if analysis['text_input_nodes']:
        print(f"\n📝 文本输入节点:")
        for i, node in enumerate(analysis['text_input_nodes'], 1):
            print(f"  {i}. 节点 {node['id']} ({node['class_type']})")
            print(f"     - 文本内容: {node['text_content']}")
    
    if analysis['connections']:
        print(f"\n🔗 节点连接:")
        for i, conn in enumerate(analysis['connections'], 1):
            print(f"  {i}. 节点 {conn['from']} → 节点 {conn['to']} (输出索引: {conn['output_index']})")
    
    if analysis['potential_issues']:
        print(f"\n⚠️  潜在问题:")
        for i, issue in enumerate(analysis['potential_issues'], 1):
            print(f"  {i}. {issue}")

def print_simulation_results(results: Dict[str, Any]) -> None:
    """打印模拟执行结果"""
    print("\n" + "="*60)
    print("🚀 模拟执行结果")
    print("="*60)
    
    for node_id, result in results.items():
        print(f"\n📦 节点 {node_id}:")
        
        if result["status"] == "success":
            output = result["output"]
            print(f"  ✅ 执行成功")
            print(f"  📝 文本内容: '{output['text']}'")
            print(f"  📁 文件数量: {len(output['files'])}")
            print(f"  📊 JSON数据: {len(output['json'])} 项")
            
            if "metadata" in output:
                print(f"  🏷️  包含元数据: 是")
            
            print(f"\n  🔍 完整输出:")
            print("  " + "\n  ".join(result["formatted_output"].split("\n")))
            
        else:
            print(f"  ❌ 执行失败: {result['error']}")

def main():
    """主函数"""
    print("ComfyUI Dify Integration - 工作流调试器")
    print("="*60)
    
    # 检查命令行参数
    if len(sys.argv) < 2:
        print("使用方法: python debug_workflow.py <workflow.json>")
        print("\n可用的测试工作流:")
        print("  - simple_test_workflow.json (推荐)")
        print("  - test_workflow.json")
        print("\n或提供你自己的工作流文件路径")
        return
    
    workflow_path = sys.argv[1]
    
    # 加载工作流
    workflow = load_workflow(workflow_path)
    if not workflow:
        return
    
    print(f"✅ 成功加载工作流: {workflow_path}")
    
    # 分析工作流
    analysis = analyze_workflow(workflow)
    print_analysis(analysis)
    
    # 模拟执行
    simulation_results = simulate_dify_output(workflow)
    if simulation_results:
        print_simulation_results(simulation_results)
    else:
        print("\n⚠️  未找到Dify输出节点，无法模拟执行")
    
    # 提供建议
    print("\n" + "="*60)
    print("💡 调试建议")
    print("="*60)
    
    if not analysis['dify_output_nodes']:
        print("❌ 未找到DifyTextOutputNode节点")
        print("   建议: 确保工作流中包含DifyTextOutputNode节点")
    
    if analysis['potential_issues']:
        print("⚠️  发现潜在问题，请检查上述问题列表")
    
    for node in analysis['dify_output_nodes']:
        if isinstance(node['inputs'].get('text'), list):
            source_id = node['inputs']['text'][0]
            if source_id in workflow:
                source_node = workflow[source_id]
                source_text = source_node.get('inputs', {}).get('text', '')
                if not source_text:
                    print(f"⚠️  节点 {source_id} 的文本内容为空")
                    print(f"   建议: 检查源节点 {source_id} 是否正确输出文本")
    
    print("\n如果问题仍然存在，请:")
    print("1. 确保已正确安装ComfyUI Dify Integration插件")
    print("2. 检查ComfyUI控制台日志")
    print("3. 运行 python debug_helper.py 进行完整诊断")

if __name__ == "__main__":
    main() 